import os
import cv2
import numpy as np
import matplotlib.pyplot as plt


def plot_img(image, pts, linewidth=1.5):
    fig = plt.figure()
    ax = fig.add_subplot(1, 1, 1)
    ax.imshow(image.astype(np.uint8))
    x1, y1 = min(pts[:,0]), min(pts[:, 1])
    x2, y2 = max(pts[:,0]), max(pts[:, 1])
    w, h = x2-x1, y2-y1

    for x,y in pts:
        circle = plt.Circle((x,y),2, color='r', fill=False)
        ax.add_patch(circle)

    rect = plt.Rectangle((x1, y1), w, h, fill=False, edgecolor='r', linewidth=linewidth)
    ax.add_patch(rect)

    plt.axis('off')
    plt.tight_layout()
    plt.draw()
    plt.show()
    # plt.savefig('imgs/%f.jpg'%i)

def plot_img_bbox_landm(image, bboxes, landms, linewidth=1.5):
    fig = plt.figure()
    ax = fig.add_subplot(1, 1, 1)
    ax.imshow(image[:, :, ::-1])
    for box, landm in zip(bboxes, landms):
        x1, y1, x2, y2 = box 
        rect = plt.Rectangle((x1, y1), x2-x1, y2-y1, fill=False, edgecolor='r', linewidth=1)
        ax.add_patch(rect)
        for x,y in landm.reshape(5,2)[[0, 2, 3]]:
            circle = plt.Circle((x,y), 1 , color='r', fill=False)
            ax.add_patch(circle)
        for x,y in landm.reshape(5,2)[[1, 4]]:
            circle = plt.Circle((x,y), 1 , color='b', fill=False)
            ax.add_patch(circle)
    plt.draw()
    plt.show()


def cv_plot_img_bbox_landm(image, bboxes, landms, save_path=None):
    for box, landm in zip(bboxes, landms):
        cv2.rectangle(image, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), (2, 255, 0), 1)
        for i in range(0, 5):
            cv2.circle(image, (int(landm[i * 2]), int(landm[i * 2 + 1])), 2, (0, 0, 255), -1)
    if save_path is not None:
        cv2.imwrite(save_path, image)
    else:
        cv2.imshow('demo', image)
        cv2.waitKey(0)